package com.hmtech.ayp.openai;

import cn.hutool.core.codec.Base64Decoder;
import cn.hutool.core.io.IoUtil;
import com.alibaba.fastjson2.JSON;
import lombok.extern.slf4j.Slf4j;
import org.apache.http.client.config.RequestConfig;
import org.apache.http.client.methods.CloseableHttpResponse;
import org.apache.http.client.methods.HttpPost;
import org.apache.http.entity.ContentType;
import org.apache.http.entity.mime.HttpMultipartMode;
import org.apache.http.entity.mime.MultipartEntityBuilder;
import org.apache.http.entity.mime.content.FileBody;
import org.apache.http.entity.mime.content.StringBody;
import org.apache.http.impl.client.CloseableHttpClient;
import org.apache.http.impl.client.HttpClients;
import org.springframework.stereotype.Component;
import org.springframework.web.multipart.MultipartFile;

import java.io.File;
import java.io.FileOutputStream;

@Component
@Slf4j
public class OpenAIClient {
    private final String apiKey = "sk-aBSir7qCZhRgpzP4E0722aC33c5247F4877eB9746fEe0729";
    private final String baseUrl = "https://api.laozhang.ai/v1";

    private final CloseableHttpClient client = HttpClients.custom()
            .build();

    public String transfer(String prompt, MultipartFile mpFile) throws Exception {
        // 1) Build a RequestConfig with 2-minute timeouts
        RequestConfig config = RequestConfig.custom()
                // Maximum time to establish the connection
                .setConnectTimeout(120_000)
                // Maximum time waiting for a connection from the pool
                .setConnectionRequestTimeout(120_000)
                // Maximum time waiting for data (a “socket” timeout)
                .setSocketTimeout(120_000)
                .build();
        HttpPost post = new HttpPost(baseUrl + "/images/edits");
        post.setConfig(config);
        post.setHeader("Authorization", "Bearer " + apiKey);
        // 表单参数列表
        MultipartEntityBuilder builder = MultipartEntityBuilder.create();
        builder.addPart("model", new StringBody("gpt-image-1", ContentType.TEXT_PLAIN));
        builder.addPart("quality", new StringBody("low", ContentType.TEXT_PLAIN));
        builder.addPart("prompt", new StringBody(prompt, ContentType.TEXT_PLAIN));
        builder.setMode(HttpMultipartMode.BROWSER_COMPATIBLE);
        File tempFile = new File(System.getProperty("java.io.tmpdir") + "/" + mpFile.getOriginalFilename());
        mpFile.transferTo(tempFile);
        builder.addPart("image", new FileBody(tempFile, ContentType.DEFAULT_BINARY, mpFile.getName()));
        post.setEntity(builder.build());
        CloseableHttpResponse response = client.execute(post);
        int code = response.getStatusLine().getStatusCode();
        byte[] bytes = response.getEntity().getContent().readAllBytes();
        if (code == 200) {
            String imgBase64 = JSON.parseObject(new String(bytes), ImagesResponse.class).getData().get(0).getB64Json();
            saveTmpFile(imgBase64);
            return imgBase64;
        } else {
            log.error("code:{},msg:{},content:{}", code, response.getStatusLine().getReasonPhrase(), new String(bytes));
        }
        return null;
    }

    private void saveTmpFile(String imgBase64) throws Exception {
        String fileName = "./tmp/" + System.currentTimeMillis() + ".jpg";
        File file = new File(fileName);
        IoUtil.write(new FileOutputStream(file), true, Base64Decoder.decode(imgBase64));
    }
}
